package com.example.service;

import com.example.data.entity.CleanedData;
import com.example.data.entity.PreprocessedData;
import com.example.data.entity.kg.KgEntity;
import com.example.data.entity.kg.KgRelation;
import com.example.data.mapper.CleanedDataMapper;
import com.example.data.mapper.KgEntityMapper;
import com.example.data.mapper.KgRelationMapper;
import com.example.data.mapper.PreprocessedDataMapper;
import com.example.util.LlmClientUtil;
import com.google.gson.Gson;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import jakarta.annotation.Resource;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.Set;
import java.util.HashSet;

/**
 * 关系抽取服务：从预处理数据中提取实体间的关系，存入kg_relation表
 * 依赖已抽取的实体（kg_entity表）和预处理文本（preprocessed_data表）
 */
@Service
public class RelationExtractService {
    private static final Logger log = LoggerFactory.getLogger(RelationExtractService.class);
    private final Gson gson = new Gson();

    // 支持的关系类型（与kg_relation表的relation_type字段对应，避免LLM生成无效类型）
    private static final List<String> SUPPORTED_RELATION_TYPES =
            List.of("定都", "建立", "统治", "任用", "发生于", "属于", "击败", "修建");

    // 关系抽取Prompt模板：强制LLM输出结构化CSV，关联已存实体
    private static final String RELATION_EXTRACT_PROMPT = """
            任务：从以下秦朝史料文本中提取实体间的关系，严格遵循约束规则。
            
            一、已知实体列表（仅可使用这些实体，格式：名称|类型）：
            %s
            
            二、允许的关系类型（仅可使用以下类型）：
            %s
            
            三、输出格式（必须为CSV，无表头，字段用英文逗号分隔）：
            主体实体名称,关系类型,客体实体名称,关系描述,置信度（0-1，保留4位小数）
            
            四、约束规则：
            1. 主体和客体必须在“已知实体列表”中（名称完全匹配，不允许新增实体）；
            2. 关系类型必须在“允许的关系类型”中，不允许自定义；
            3. 关系描述需从文本中提取具体依据（如“嬴政定都咸阳：文本提及‘嬴政统一后定都咸阳’”）；
            4. 置信度基于文本相关性打分（明确提及→0.9+，间接推断→0.6-0.8，模糊关联→0.3-0.5）；
            5. 排除重复关系（同一主体+同一关系+同一客体视为重复）。
            
            五、待处理文本：
            %s
            """;

    private void markRelationExtracted(Long preprocessedId, Integer status) {
        try {
            preprocessedDataMapper.markRelationExtracted(preprocessedId, status);
        } catch (Exception e) {
            log.error("标记关系抽取状态失败，preprocessed_id={}, status={}", preprocessedId, status, e);
        }
    }

    /**
     * 关系去重：基于“主体ID+客体ID+关系类型”判断重复
     * 避免向数据库插入重复关系，减轻数据库唯一索引的压力
     *
     * @param relations 待去重的关系列表
     * @return 去重后的关系列表
     */
    private List<KgRelation> deduplicateRelations(List<KgRelation> relations) {
        // 使用Set存储唯一键，实现去重
        Set<String> uniqueKeys = new HashSet<>();
        List<KgRelation> uniqueRelations = new ArrayList<>();

        for (KgRelation relation : relations) {
            // 构建唯一键：主体ID+客体ID+关系类型（与kg_relation表的唯一索引一致）
            String uniqueKey = buildRelationUniqueKey(relation);

            // 若该键不存在，则添加到去重列表和键集合中
            if (!uniqueKeys.contains(uniqueKey)) {
                uniqueKeys.add(uniqueKey);
                uniqueRelations.add(relation);
            } else {
                log.debug("发现重复关系，已过滤：主体ID={}, 客体ID={}, 关系类型={}",
                        relation.getSubjectEntityId(),
                        relation.getObjectEntityId(),
                        relation.getRelationType());
            }
        }

        log.debug("关系去重完成，原数量={}, 去重后数量={}", relations.size(), uniqueRelations.size());
        return uniqueRelations;
    }

    /**
     * 构建关系唯一键：用于判断两个关系是否重复
     *
     * @param relation 关系对象
     * @return 唯一键字符串（格式：subjectId_objectId_relationType）
     */
    private String buildRelationUniqueKey(KgRelation relation) {
        return String.format("%d_%d_%s",
                relation.getSubjectEntityId(),
                relation.getObjectEntityId(),
                relation.getRelationType());
    }

    @Resource
    private PreprocessedDataMapper preprocessedDataMapper;

    @Resource
    private KgEntityMapper kgEntityMapper;

    @Resource
    private KgRelationMapper kgRelationMapper;

    @Resource
    private CleanedDataMapper cleanedDataMapper;

    @Resource
    private LlmClientUtil llmClientUtil;

    /**
     * 批量抽取关系：查询未抽取关系的预处理数据，批量处理
     *
     * @param batchSize 单次处理数量（避免内存溢出，建议20-50）
     * @return 成功抽取的关系总数
     */
    @Transactional(rollbackFor = Exception.class)
    public int batchExtractRelations(Integer batchSize) {
        log.info("开始批量抽取关系，批次大小：{}", batchSize);

        // 1. 加载已存实体：构建“实体名称→实体对象”映射（用于快速匹配ID）
        Map<String, KgEntity> entityMap = loadValidEntityMap();
        if (entityMap.isEmpty()) {
            log.error("kg_entity表中无有效实体，无法抽取关系");
            return 0;
        }

        // 2. 查询未抽取关系的预处理数据（需先在preprocessed_data表新增is_relation_extracted字段）
        List<PreprocessedData> unprocessedList = preprocessedDataMapper.selectUnExtractedRelations(batchSize);
        if (unprocessedList.isEmpty()) {
            log.info("无未抽取关系的预处理数据");
            return 0;
        }

        // 3. 逐条处理预处理数据
        int totalSuccess = 0;
        for (PreprocessedData preData : unprocessedList) {
            try {
                log.info("处理预处理数据，ID：{}", preData.getId());

                // 4. 校验该数据是否已抽取过关系（避免重复处理）
                List<KgRelation> existRelations = kgRelationMapper.selectByPreprocessedId(preData.getId());
                if (!existRelations.isEmpty()) {
                    log.debug("该预处理数据已抽取关系（数量：{}），跳过", existRelations.size());
                    markRelationExtracted(preData.getId(), 1); // 标记为“已抽取”
                    continue;
                }

                // 5. 构建关系抽取Prompt
                String prompt = buildRelationPrompt(preData, entityMap);

                // 6. 调用LLM获取关系结果
                String llmResponse = llmClientUtil.callLlm(prompt);
                if (llmResponse == null || llmResponse.trim().isEmpty() || llmResponse.contains("无")) {
                    log.debug("该文本无关系可抽取，预处理ID：{}", preData.getId());
                    markRelationExtracted(preData.getId(), 1); // 标记为“已抽取（无结果）”
                    continue;
                }

                // 7. 解析LLM结果为KgRelation列表
                List<KgRelation> relationList = parseLlmResponse(llmResponse, preData, entityMap);
                if (relationList.isEmpty()) {
                    log.warn("关系解析失败，预处理ID：{}", preData.getId());
                    markRelationExtracted(preData.getId(), 2); // 标记为“抽取失败”
                    continue;
                }

                // 8. 代码层去重（减少数据库压力）
                List<KgRelation> uniqueRelations = deduplicateRelations(relationList);

                // 9. 批量入库
                kgRelationMapper.batchInsert(uniqueRelations);
                totalSuccess += uniqueRelations.size();
                log.info("预处理ID：{}，成功抽取关系：{} 条", preData.getId(), uniqueRelations.size());

                // 10. 标记预处理数据为“已抽取关系”
                markRelationExtracted(preData.getId(), 1);

            } catch (Exception e) {
                log.error("处理预处理数据ID：{} 失败", preData.getId(), e);
                markRelationExtracted(preData.getId(), 2); // 标记为“抽取失败”
            }
        }

        log.info("批量关系抽取完成，共处理预处理数据：{} 条，成功抽取关系：{} 条", unprocessedList.size(), totalSuccess);
        return totalSuccess;
    }

    /**
     * 加载已存有效实体：从kg_entity表查询is_valid=1的实体，构建名称→实体映射
     */
    private Map<String, KgEntity> loadValidEntityMap() {
        List<KgEntity> entityList = kgEntityMapper.selectAll();
        Map<String, KgEntity> entityMap = new HashMap<>();

        log.info("从kg_entity表查询到的总实体数：{}", entityList.size()); // 打印查询总数

        for (KgEntity entity : entityList) {
            // 打印每个实体的关键信息（便于排查）
            log.debug("实体ID：{}，名称：{}，is_valid：{}，类型：{}",
                    entity.getId(), entity.getEntityName(), entity.getIsValid(), entity.getEntityType());

            // 仅保留有效实体
            if (entity.getIsValid() != null && entity.getIsValid() == 1) {
                // 检查实体名称是否为空
                if (entity.getEntityName() == null || entity.getEntityName().trim().isEmpty()) {
                    log.warn("实体ID：{} 的名称为空，跳过存入映射表", entity.getId());
                    continue;
                }
                // 存入实体名称→实体对象
                entityMap.put(entity.getEntityName().trim(), entity);
                // 处理别名
                if (entity.getEntityAlias() != null && !entity.getEntityAlias().isEmpty()) {
                    String[] aliases = entity.getEntityAlias().split("，");
                    for (String alias : aliases) {
                        alias = alias.trim();
                        if (!alias.isEmpty()) {
                            entityMap.put(alias, entity);
                            log.debug("实体ID：{} 的别名：{} 已存入映射表", entity.getId(), alias);
                        }
                    }
                }
            } else {
                log.debug("实体ID：{} 无效（is_valid={}），跳过存入映射表",
                        entity.getId(), entity.getIsValid());
            }
        }

        log.info("实体映射表最终大小：{}，包含的实体名称：{}",
                entityMap.size(), entityMap.keySet()); // 打印映射表的所有键（实体名称）
        return entityMap;
    }

    /**
     * 构建关系抽取Prompt：关联已存实体列表和预处理文本
     */
    private String buildRelationPrompt(PreprocessedData preData, Map<String, KgEntity> entityMap) {
        // 1. 构建已知实体列表（格式：名称|类型）
        String entityListStr = entityMap.values().stream()
                .map(entity -> entity.getEntityName() + "|" + entity.getEntityType())
                .collect(Collectors.joining("\n"));

        // 2. 构建允许的关系类型字符串
        String supportedRelationsStr = String.join("、", SUPPORTED_RELATION_TYPES);

        // 3. 准备输入文本（优先用实体预标注文本，降低LLM识别难度）
        String inputText = preData.getAnnotatedText() != null ?
                preData.getAnnotatedText() :
                gson.fromJson(preData.getSentences(), List.class).toString();

        // 4. 填充Prompt模板
        return String.format(RELATION_EXTRACT_PROMPT, entityListStr, supportedRelationsStr, inputText);
    }



    /**
     * 解析LLM返回的CSV结果，转化为KgRelation对象
     */
    private List<KgRelation> parseLlmResponse(
            String llmResponse,
            PreprocessedData preData,
            Map<String, KgEntity> entityMap) {

        List<KgRelation> relationList = new ArrayList<>();
        String[] lines = llmResponse.split("\\r?\\n"); // 处理Windows/Linux换行符

        for (String line : lines) {
            line = line.trim();
            // 跳过空行或表头（若LLM误输出表头）
            if (line.isEmpty() || line.startsWith("主体实体名称")) {
                continue;
            }

            // 处理关系描述中含逗号的情况：用最后一个逗号分割置信度
            int lastCommaIndex = line.lastIndexOf(",");
            if (lastCommaIndex == -1 || lastCommaIndex == line.length() - 1) {
                log.warn("无效关系行（字段分割失败）：{}", line);
                continue;
            }

            // 提取字段
            String basePart = line.substring(0, lastCommaIndex).trim(); // 主体,关系,客体,描述
            String confidenceStr = line.substring(lastCommaIndex + 1).trim(); // 置信度
            String[] baseFields = basePart.split(",", 4); // 最多分割4段（适配描述含逗号）

            // 校验字段数量（必须为4段）
            if (baseFields.length != 4) {
                log.warn("无效关系行（字段数量错误）：{}", line);
                continue;
            }

            // 清洗字段
            String subjectName = baseFields[0].trim();    // 主体实体名称
            String relationType = baseFields[1].trim();   // 关系类型
            String objectName = baseFields[2].trim();     // 客体实体名称
            String relationDesc = baseFields[3].trim();   // 关系描述

            // -------------- 关键校验：确保数据合法 --------------
            // 1. 校验关系类型是否支持
            if (!SUPPORTED_RELATION_TYPES.contains(relationType)) {
                log.warn("无效关系类型（不支持）：{}，行：{}", relationType, line);
                continue;
            }

            // 2. 从实体映射表中匹配主体/客体的实体对象（确保实体已存在于kg_entity表）
            KgEntity subjectEntity = entityMap.get(subjectName);
            KgEntity objectEntity = entityMap.get(objectName);
            if (subjectEntity == null) {
                log.warn("主体实体不存在（未在kg_entity表中）：{}，行：{}", subjectName, line);
                continue;
            }
            if (objectEntity == null) {
                log.warn("客体实体不存在（未在kg_entity表中）：{}，行：{}", objectName, line);
                continue;
            }

            // 3. 校验置信度格式（必须是0-1的小数，保留4位）
            Double confidence;
            try {
                confidence = Double.parseDouble(confidenceStr);
                // 校验范围（0≤置信度≤1）
                if (confidence < 0 || confidence > 1) {
                    throw new IllegalArgumentException("置信度超出0-1范围");
                }
                // 保留4位小数（避免数据库decimal(5,4)字段溢出）
                confidence = Math.round(confidence * 10000.0) / 10000.0;
            } catch (Exception e) {
                log.warn("置信度格式错误：{}，行：{}，错误：{}", confidenceStr, line, e.getMessage());
                continue;
            }

            // 4. 校验关系描述非空（避免无意义描述）
            if (relationDesc.isEmpty() || relationDesc.length() < 10) {
                log.warn("关系描述无效（过短或空）：{}，行：{}", relationDesc, line);
                continue;
            }

            // -------------- 构建KgRelation对象（与表字段一一对应） --------------
            KgRelation relation = new KgRelation();
            // 主体/客体实体ID（关联kg_entity表的主键）
            relation.setSubjectEntityId(subjectEntity.getId());
            relation.setObjectEntityId(objectEntity.getId());
            // 关系核心信息
            relation.setRelationType(relationType);
            relation.setRelationDesc(relationDesc);
            relation.setConfidence(confidence);
            // 来源链路（关联预处理→清洗→原始数据，确保可追溯）
            relation.setSourcePreprocessedId(preData.getId()); // 来源预处理数据ID
            // 从预处理数据关联清洗数据ID（preprocessed_data.cleaned_id）
            relation.setSourceCleanedId(preData.getCleanedId());
            // 从清洗数据关联原始数据ID（需查询cleaned_data表的raw_id字段）
            CleanedData cleanedData = cleanedDataMapper.selectById(preData.getCleanedId());
            if (cleanedData != null) {
                relation.setSourceRawId(cleanedData.getRawId());
            }
            // 基础状态字段
            relation.setIsValid(1); // 默认为有效（1=有效，0=无效）
            relation.setCreateTime(LocalDateTime.now()); // 创建时间
            relation.setUpdateTime(LocalDateTime.now()); // 更新时间

            // 将合法的关系添加到列表
            relationList.add(relation);
        }

        log.debug("LLM结果解析完成，有效关系数量：{}", relationList.size());
        return relationList;
    }
}